import cv2
import numpy as np
from ultralytics import YOLO
from tkinter import Tk, Button, filedialog, Label, Entry, ttk, messagebox
from PIL import Image, ImageDraw, ImageFont
import threading
import torch
import os
import concurrent.futures
import hashlib
from time import time

# 设置最大线程数
max_threads = os.cpu_count()
torch.set_num_threads(max_threads)
os.environ["OMP_NUM_THREADS"] = str(max_threads)

# 设置CUDA优化
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False

# 创建GUI窗口
root = Tk()
root.title("YOLO图片处理程序")
root.geometry("600x700")
root.resizable(False, False)

# 设置背景色和字体
root.configure(bg="#f0f0f0")

style = ttk.Style()
style.configure("TButton", font=("Arial", 12), width=20, padding=10)
style.configure("TLabel", font=("Arial", 12), padding=5)
style.configure("TEntry", font=("Arial", 12), width=30)
style.configure("TProgressbar", thickness=30)

# 全局变量
image_dir = None
model_path = None
output_dir = None
labels_dir = None
target_width = 416
target_height = 416

# 创建进度条
progress = ttk.Progressbar(root, length=400, mode='determinate', maximum=100)
progress.pack(side="bottom", pady=10)

# 加载YOLO模型
def load_model():
    global model_path
    model_path = filedialog.askopenfilename(title="选择模型文件", filetypes=[("PT files", "*.pt")])
    if model_path:
        model_label.config(text=f"选择的模型: {model_path.split('/')[-1]}")
    else:
        model_label.config(text="未选择模型")

# 选择图片目录
def choose_image_dir():
    global image_dir
    image_dir = filedialog.askdirectory(title="选择图片目录")
    if image_dir:
        image_dir_label.config(text=f"选择的目录: {image_dir.split('/')[-1]}")
    else:
        image_dir_label.config(text="未选择图片目录")

# 选择输出目录
def choose_output_dir():
    global output_dir
    output_dir = filedialog.askdirectory(title="选择结果图保存目录")
    if output_dir:
        output_label.config(text=f"保存结果图的目录: {output_dir}")
    else:
        output_label.config(text="未选择保存目录")

# 选择标记信息保存目录
def choose_labels_dir():
    global labels_dir
    labels_dir = filedialog.askdirectory(title="选择标记信息保存目录")
    if labels_dir:
        labels_label.config(text=f"保存标记信息的目录: {labels_dir}")
    else:
        labels_label.config(text="未选择标记信息目录")

# 生成基于类别标签的颜色
def generate_color(label):
    hash_object = hashlib.md5(label.encode())
    hash_value = hash_object.hexdigest()

    r = int(hash_value[:2], 16)
    g = int(hash_value[2:4], 16)
    b = int(hash_value[4:6], 16)

    r = min(255, max(180, r + 50))  
    g = min(255, max(180, g + 50))  
    b = min(255, max(180, b + 50))  

    return (r, g, b)

# 并行处理每一张图片
def process_image(image_path, model, device, target_width, target_height, output_dir, labels_dir, progress_callback):
    image = cv2.imread(image_path)
    original_image = image.copy()

    # 使用 YOLO 进行对象检测
    results = model(image)

    # 获取检测框和标签信息
    boxes = results[0].boxes
    labels = results[0].names  # 分类标签
    confs = boxes.conf  # 检测置信度

    # 保存YOLO格式标记信息
    label_file = os.path.join(labels_dir, f"{os.path.splitext(os.path.basename(image_path))[0]}.txt")
    with open(label_file, "w") as f:
        for i, box in enumerate(boxes.xyxy):  
            x1, y1, x2, y2 = box
            class_id = int(boxes.cls[i])
            label = labels[class_id]
            conf = confs[i]

            # YOLO格式：class_id x_center y_center width height
            img_width, img_height = image.shape[1], image.shape[0]
            x_center = (x1 + x2) / 2 / img_width
            y_center = (y1 + y2) / 2 / img_height
            width = (x2 - x1) / img_width
            height = (y2 - y1) / img_height

            f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

    # 绘制检测框
    for i, box in enumerate(boxes.xyxy):
        x1, y1, x2, y2 = box
        color = generate_color(labels[int(boxes.cls[i])])
        cv2.rectangle(original_image, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
        label = f"{labels[int(boxes.cls[i])]} {confs[i]:.2f}"

        # 计算文本大小
        (text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

        # 初步设置文本位置
        text_x = int(x1)
        text_y = int(y1) - 30  # 初步设置文本在框的上方

        # 调整文本位置，确保文本不超出边界
        if text_x < 0:
            text_x = 0

        if text_y < text_height:
            text_y = int(y2) + text_height + 5  # 让文本位于框的下方

        if text_x + text_width > image.shape[1]:
            text_x = image.shape[1] - text_width

        if text_y + text_height > image.shape[0]:
            text_y = int(y1) - text_height - 5  # 让文本位于框的上方

        # 绘制文本
        original_image = draw_text_with_chinese(original_image, label, (text_x, text_y), color)

    # 保存结果图
    result_image_path = os.path.join(output_dir, os.path.basename(image_path))
    cv2.imwrite(result_image_path, original_image)

    progress_callback()  # 更新进度条


# 使用Pillow加载字体
def draw_text_with_chinese(frame, text, position, color, font_path="C:\\Windows\\Fonts\\simhei.ttf", font_size=20):
    pil_img = Image.fromarray(frame)
    draw = ImageDraw.Draw(pil_img)

    font = ImageFont.truetype(font_path, font_size)
    x, y = position
    draw.text((x, y), text, font=font, fill=color)

    return np.array(pil_img)

# 更新进度条
def update_progress(progress_value):
    progress['value'] = progress_value
    root.update_idletasks()

# 执行图片处理任务（此函数将用于线程）
def process_images_task():
    global target_width, target_height
    target_width = int(target_width_entry.get())  
    target_height = int(target_height_entry.get())  

    if not image_dir or not model_path or not output_dir or not labels_dir:
        status_label.config(text="请选择图片目录、模型和保存路径")
        return

    # 检测是否支持 CUDA
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # 在主线程中更新设备标签
    root.after(0, update_device_label, device)

    # 每个线程加载自己的模型
    model = YOLO(model_path).to(device)
    model.eval()  
    model.fuse()  

    # 获取图片目录下所有图片
    image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

    total_images = len(image_paths)
    progress_step = 100 / total_images if total_images > 0 else 100

    # 使用线程池并行处理每一张图片
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
        futures = []
        for i, image_path in enumerate(image_paths):
            futures.append(executor.submit(process_image, image_path, model, device, target_width, target_height, output_dir, labels_dir, lambda i=i: update_progress((i + 1) * progress_step)))

        # 等待所有任务完成
        for future in futures:
            future.result()

    status_label.config(text="处理完成")
    messagebox.showinfo("完成", "所有图片已处理完毕！")

# 更新设备标签的函数
def update_device_label(device):
    device_label.config(text=f"当前使用设备: {device.upper()}")

# 开始处理按钮
def start_processing():
    status_label.config(text="处理中，请稍候...")
    threading.Thread(target=process_images_task).start()

# 选择目标宽度和高度
target_width_label = ttk.Label(root, text="目标宽度")
target_width_label.pack(pady=5)

target_width_entry = ttk.Entry(root)
target_width_entry.insert(0, "416")  
target_width_entry.pack(pady=5)

target_height_label = ttk.Label(root, text="目标高度")
target_height_label.pack(pady=5)

target_height_entry = ttk.Entry(root)
target_height_entry.insert(0, "416")  
target_height_entry.pack(pady=5)

model_button = ttk.Button(root, text="选择YOLO模型", command=load_model)
model_button.pack(pady=10)

model_label = ttk.Label(root, text="未选择模型")
model_label.pack(pady=5)

image_dir_button = ttk.Button(root, text="选择图片目录", command=choose_image_dir)
image_dir_button.pack(pady=10)

image_dir_label = ttk.Label(root, text="未选择图片目录")
image_dir_label.pack(pady=5)

output_button = ttk.Button(root, text="选择结果图保存目录", command=choose_output_dir)
output_button.pack(pady=10)

output_label = ttk.Label(root, text="未选择保存目录")
output_label.pack(pady=5)

labels_button = ttk.Button(root, text="选择标记信息保存目录", command=choose_labels_dir)
labels_button.pack(pady=10)

labels_label = ttk.Label(root, text="未选择标记信息目录")
labels_label.pack(pady=5)

# 启动处理按钮
process_button = ttk.Button(root, text="开始处理", command=start_processing)
process_button.pack(pady=20)

# 状态信息标签
status_label = ttk.Label(root, text="", foreground="red")
status_label.pack(pady=5)

# 设备信息标签
device_label = ttk.Label(root, text="当前使用设备: ")
device_label.pack(pady=5)

# 启动GUI
root.mainloop()
